import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict,namedtuple
import seaborn as sns
import pandas as pd

class BsAbs():
    def __init__(self,koncd3,koffcd3,kontaa,kofftaa) -> None:
        self.koncd3 = koncd3
        self.koffcd3 = koffcd3
        self.kontaa = kontaa
        self.kofftaa = kofftaa

class Cell():
    def __init__(self, taa_e=0, cd3_e=0,rep_rate=0) -> None:
        '''
        dose: n/L 细胞浓度
        taa_e: taa表达量
        cd3_e: cd3表达量
        '''
        self.taa_e = taa_e # n per cell
        self.cd3_e = cd3_e # n percell
        self.rep_rate = rep_rate

    @property
    def cd3_rate(self):
        '''cd3: nmol/ cell
        '''
        return (self.cd3_e/6.023e+14)
    
    @property
    def taa_rate(self):
        '''taa: nmol / cell
        '''
        return (self.taa_e/6.023e+14)
    
class TumorEmacEffect():
    def __init__(self,volumn,kmax,kc50,kg0,kg,tau,Mmax,phi) -> None:
        '''
        volumn 肿瘤初始体积
        kmax 最大杀伤率
        kc50 半杀伤率浓度
        kg0 指数增长率
        kg 线性增长率
        tau 肿瘤各部位的转移时间
        Mmax 最大肿瘤体积
        '''
        self.m1 = volumn
        self.m2 = 0
        self.m3 = 0
        self.m4 = 0
        self.kmax = kmax
        self.kc50 = kc50
        self.kg0 = kg0
        self.kg = kg
        self.tau = tau
        self.Mmax = Mmax
        self.phi = phi

    @property
    def w(self):
        return self.m1 + self.m2 + self.m3 + self.m4
    
    @property
    def delta(self,Trimer):
        kill = (self.kmax+Trimer)/(self.kc50+Trimer)
        delta = {}
        delta['m1'] = (self.kg0 *(1-self.totalvolume/self.Mmax) /((1+(self.kg0/self.kg * self.w)**self.phi)**(1/self.phi)) - kill)* self.m1
        delta['m2'] = kill * self.m1 -self.m2/self.tau
        delta['m3'] = (self.m2-self.m3)/self.tau
        delta['m4'] = (self.m3-self.m4)/self.tau
        return delta

    def update(self):
        for k,v in self.delta.items():
            eval(f"self.{k} += {v}")


class Compartment():
    def __init__(self,name,volumn,drug,tcell,tumor,out_rate:dict,clear_rate:dict,rep_rate:dict,Drug: BsAbs,Tcell: Cell,Tumor: Cell) -> None:
        '''
        Tcell,Tumor,Drug: 细胞和药物的参数
        name: 室的名字
        cd3: CD3量
        cd3d: CD3-drug 二聚体结合量
        taad: taa-drug 二聚体结合量
        taacd3d: cd3-taa-drug 三聚体结合梁
        drug: 药物剂量
        taa: TAA量
        volumn: 室容积
        out_rate: 输出率，输出到其他部位的比例
        clear_rate: 药物清除率
        '''
        self.name = name
        self.volumn = volumn
        self.out_rate = out_rate
        self.clear_rate = clear_rate
        self.rep_rate = rep_rate
        self.Drug = Drug
        self.Tcell = Tcell
        self.Tumor = Tumor
        self.init(drug,tcell,tumor)

    
    def init(self,drug,tcell,tumor):
        '''
        abs 药物剂量nmol/L
        tcell 细胞数量
        tumor 细胞数量
        '''
        self.timedatas = defaultdict(list)
        self.timedatas['times'] = [0]
        self.timedatas['drug'] = [drug*self.volumn]
        self.tcell = tcell
        self.tumor = tumor
        cd3 = tcell * self.Tcell.cd3_rate # nmol
        taa = tumor * self.Tumor.taa_rate # nmol
        # print(self.Tumor.taa_rate)
        # print(tcell,cd3,tumor,taa)
        self.timedatas['cd3'] = [cd3]
        self.timedatas['taa'] = [taa]
        self.timedatas['cd3d'] = [0]
        self.timedatas['taad'] = [0]
        self.timedatas['taacd3d'] = [0]


    # 先写属性
    @property
    def drug(self):
        return self.timedatas['drug'][-1]
    
    @property
    def cdrug(self):
        return self.drug/self.volumn

    @property
    def taa(self):
        return self.timedatas['taa'][-1]

    @property
    def ctaa(self):
        return self.taa/self.volumn

    @property
    def cd3(self):
        return self.timedatas['cd3'][-1]
    @property
    def ccd3(self):
        return self.cd3/self.volumn

    @property
    def taad(self):
        return self.timedatas['taad'][-1]
    
    @property
    def ctaad(self):
        return self.taad/self.volumn

    @property
    def cd3d(self):
        return self.timedatas['cd3d'][-1]

    @property
    def ccd3d(self):
        return self.cd3d/self.volumn
    
    @property
    def taacd3d(self):
        return self.timedatas['taacd3d'][-1]

    @property
    def ctaacd3d(self):
        return self.taacd3d/self.volumn
    
    
    # recorder 记录实验数据变化
    @property
    def metric(self):
        table = pd.DataFrame.from_dict(self.timedatas,dtype=float)
        return table.set_index('times')

    # action 部分
    @property    
    def clear(self):
        clear = defaultdict(float)
        if self.clear_rate is not None:
            for key,value in self.clear_rate.items():
                rate = value
                clear[key] =  rate* self.timedatas[key][-1]
        return clear


    def input(self,inputs):
        self._in = defaultdict(float)
        for inp in inputs:
            if self.name in inp.keys():
                for item in inp[self.name]:
                    self._in[item] += inp[self.name][item]
 
    def out(self):
        self._out = defaultdict(float)
        tissue_out = defaultdict(dict)
        if self.out_rate is not None:
            for  key in self.out_rate:
                for item in self.out_rate[key].keys():
                    rate = self.out_rate[key][item]/self.volumn 
                    newvalue = rate* self.timedatas[item][-1]
                    tissue_out[key][item] = newvalue
                    self._out[item] += newvalue
        return tissue_out

    @property
    def copy(self):
        delta = defaultdict(float)
        if self.rep_rate:
            delta['taa'] = self.taa * self.rep_rate.get('taa') + self.taad*self.rep_rate.get('taad',0)
            delta['cd3'] = self.taa * self.rep_rate.get('cd3') + self.taad*self.rep_rate.get('cd3d',0)
        return delta

    @property
    def innerdelta(self):
        '''内部变化
        '''
        delta = defaultdict(float)
        # Monmier
        delta['drug'] = self.Drug.kofftaa * self.taad + self.Drug.koffcd3 * self.cd3d - (self.Drug.kontaa * self.taa  + self.Drug.koncd3 * self.cd3 ) * self.drug/self.volumn 
        delta['taa'] = (self.taad + self.taacd3d) * self.Drug.kofftaa - (self.drug + self.cd3d) * self.Drug.kontaa * self.taa /self.volumn
        delta['cd3'] = (self.taacd3d + self.cd3d) * self.Drug.koffcd3 - (self.drug + self.taad) * self.Drug.koncd3 * self.cd3 /self.volumn
        # Dimier
        delta['taad'] = self.Drug.kontaa * self.drug * self.taa/self.volumn + self.Drug.koffcd3 * self.taacd3d - self.Drug.kofftaa * self.taad - self.Drug.koncd3 * self.taad * self.cd3/self.volumn
        delta['cd3d'] = self.Drug.koncd3 * self.drug * self.cd3/self.volumn + self.Drug.kofftaa * self.taacd3d - self.Drug.koffcd3 * self.cd3d - self.Drug.kontaa * self.cd3d * self.taa/self.volumn
        # Trimer
        delta['taacd3d'] = self.Drug.kontaa * self.cd3d * self.taa/self.volumn + self.Drug.koncd3 * self.taad * self.cd3/self.volumn - self.taacd3d * (self.Drug.kofftaa  + self.Drug.koffcd3) 
        return delta

    @property
    def delta(self):
        delta = defaultdict(float)
        for key in self.innerdelta.keys():
            delta[key] = self._in[key] + self.innerdelta[key] + self.copy[key] - self._out[key] - self.clear[key]
        return delta
    
    def update(self,interval):
        for key,value in self.delta.items():
            last = self.timedatas[key][-1]
            self.timedatas[key].append(last+value*interval)
        self.timedatas['times'].append(self.timedatas['times'][-1]+interval)
    
    
    def show(self,key=None):
        if key:
            plt.title(f"{'-'.join(key)} in {self.name}")
            sns.lineplot(data=self.metric[key])
        else:
            sns.lineplot(data=self.metric)
        plt.show()